# -*- coding: utf-8 -*-

'''
 :author: Rayco Abad-Martín <rayco.abad@gmail.com>
 :author: Elena Marrero-Méndez <elemarez@gmail.com>
'''

from yacf.Backends.C.Files.FileTypes import H_FileType, C_FileType
from yacf.Backends.C.Visitors.OmpFilters import *
from Backends.MPI.Mutators.Common import AbstractMPIMutator, mpi_outline_loop_stmts
from yacf.Frontend.Shortcuts import getCurrentLanguageAst
from yacf.Frontend.SymbolTable import updateSymbolTableFromNode
from yacf.Tools import SourceStorage
from yacf.Tools.Tree import InsertTool, NodeNotFound, ReplaceTool, RemoveTool, getFileNode
import config

c_ast = getCurrentLanguageAst()

class MPI_OmpFor(AbstractMPIMutator):

    def __init__(self, file_name = 'loopFor', *args, **kwargs):
        self.file_name = file_name
        super(MPI_OmpFor, self).__init__(*args, **kwargs)

    def filter(self, ast):
        """ Filter definition
            Returns the first node matching with the filter
        """
        # Build a visitor , matching the OmpFor node of the AST
        f = OmpForFilter()
        node = f.apply(ast)
        self._func_def = f.get_func_def()
        self._parallel = f.get_parallel()
        return node

    def apply_all(self, parent_parallel_node, ast):
        """ Apply mutation to all matches """
        start_node = None
        self.ast = ast
        f = OmpForFilter()
        num = 0
        try:
            for elem in f.iterate(ast):
                # Save previous state
                old_file_name = self.file_name
                self.file_name = self.file_name + str(num)
                old_clauses = self._clauses
                self._clauses = {}
                # Current scope variables
                self._func_def = f.get_func_def()
                self._parallel = f.get_parallel()
                # If current node is not child of first parallel node, stop
                if self._parallel != parent_parallel_node:
                    raise StopIteration
                start_node = self.mutatorFunction(ast, elem)
                # Restore previous state
                self.file_name = old_file_name
                self._clauses = old_clauses
                num += 1
        except NodeNotFound as nf:
            print str(nf)
        except StopIteration:
            return self._parallel
        return start_node

    def mutatorFunction(self, ast, ompFor_node):

#        print "\n********************* FILTER NODE *********************\n"
#        ast.parent.show(attrnames = True)
#        print "\n******************* END FILTER NODE *******************\n"
        
        # Solve the error when the pragma parallel has not clauses
        if ompFor_node.clauses == None:
            ompFor_node.clauses = {}
        clause_dict = self._get_dict_from_clauses(ompFor_node.clauses, ast)
        reduction_params = clause_dict['REDUCTION']
        nowait = clause_dict.has_key('NOWAIT')
        # Check if nested for
        collapse = clause_dict.has_key('COLLAPSE')
        # Private declarations come from the parent parallel construct
        private_params = clause_dict['PRIVATE']
        shared_params = clause_dict['SHARED']
        lastprivate_params = clause_dict['LASTPRIVATE']
        # Capture the llc clauses
        llc_clause_dict = self._get_dict_from_llc_clauses(ompFor_node.stmt, ast)
        result_params = llc_clause_dict['RESULT']
        nc_result_params = llc_clause_dict['NC_RESULT']
        ####### begin of mutator MPI_OmpFor ########
        """ MPI mutator, writes memory transfer operations for a parallel region
        """
        # Capture the values loop for init and cond.
        loop = ompFor_node.stmt
        if result_params or nc_result_params:
            loop = loop.stmt
        # Inicialización del bucle
        loop_init = loop.init
        # Condición del bucle
        loop_cond = loop.cond
        # Actualización del bucle
        loop_next = loop.next
        # Valor inicial de la variable del bucle
        loop_initial = loop.init.rvalue
        # Valor final de la variable del bucle
        loop_finalice = str(loop.cond.right)
        if (loop.cond.op == '<'):
            loop_finalice += " - 1"
        # Variable del bucle
        loop_var = loop.init.lvalue
        # Bloque del bucle
        loop_stmt = loop.stmt
        subs_dir = dict()
        """ build the template_code """
        if ((loop_var) and (loop_initial) and (loop_finalice) and ((type(loop.next) == c_ast.UnaryOp) or (str(loop.next.rvalue) == '1'))):
            template_code = self.processLlcCodeUnit(variables = [], omp_cl = ompFor_node.clauses, ast = ast, translate = subs_dir)
        else:
            pass
        subs_dir.update({ 'loop_init': loop_init,
                           'loop_cond': loop_cond,
                           'loop_next': loop_next,
                           'loop_initial': loop_initial,
                           'loop_finalice': loop_finalice,
                           'loop_var': loop_var,
                           'loop_stmt': loop_stmt,
                        })
        if isinstance(ompFor_node.parent.parent, c_ast.Compound):
            forall_unit = self.fillTemplate(template_code, subs_dir, name = 'template forall unit', show = False)
            self.main_ast = getFileNode(ompFor_node)
            source = mpi_outline_loop_stmts(ompFor_node.stmt.stmt, file_name = self.file_name, template_source = forall_unit, ast = ast)
            genFile = C_FileType()
            self.source_storage.addFile(self.file_name, genFile)
            self.source_storage.append(self.file_name, genFile, '#include "llcomp_llc.h" \n')
            self.source_storage.append(self.file_name, genFile, source) 
        else:
            ompFor_node.parent.parent.stmt = forall_unit
        updateSymbolTableFromNode(ast)
